import argparse
import numpy as np
import healpy
from sphere.distribution import fb8
from matplotlib import pyplot as plt
import logging
logger = logging.getLogger()
logger.setLevel(level=logging.ERROR)


def fb8_to_pdf(f, i):
    print('\rprocessing event {} ...'.format(i), end='')
    return f.pdf(f.spherical_coordinates_to_nu(scan_theta, scan_ra))


parser = argparse.ArgumentParser(
        description='This program will load in data.npy and save healpix maps into FITs files under fits/')
parser.add_argument('--nside', default=512, type=int,
                    help='The healpix NSIDE setting')
parser.add_argument('--plotting', default=False, action='store_true',
                    help='Create mollview maps for each event as well')
args = parser.parse_args()

arr = np.load('data.npy')
scan_theta, scan_ra = healpy.pix2ang(args.nside, np.r_[:healpy.nside2npix(args.nside)])
scan_dec = np.pi/2 - scan_theta

comment_dict = {'id': 'IceCube HESE event ID',
                'mjd': 'Modified Julian Date',
                'ra': 'Most probable arrival right ascension [rad]',
                'dec': 'Most probable arrival declination [rad]',
                'energy': 'Most probable deposited energy [GeV]',
                'drlogl': 'Likelihood indicator of hypothesis confidence',
                'reconstruction': 'Hypothesis used for event reconstruction'}
# healpix probablity maps
for i, row in enumerate(arr):
    map = fb8_to_pdf(fb8(*row['params']), i)

    hdr = []
    for k in row.dtype.fields.keys():
        if k == 'params':
            continue
        if k == 'reconstruction':
            hdr.append(('RECO', row[k].decode('ASCII'), comment_dict[k]))
        else:
            hdr.append((k.upper(), row[k], comment_dict[k]))
    healpy.write_map(f'fits/event_{row["id"]:03d}.fits',
                     map,
                     fits_IDL=False,
                     coord='C',
                     column_names=['Probability density'],
                     column_units=['1/sr'],
                     extra_header=hdr,
                     overwrite=True)
    if args.plotting:
        healpy.mollview(map, coord='C', unit='1/sr', title=f'event{row["id"]:03d}')
        plt.show()
